import numpy as np
import scipy.special as sc
from chaco.api import ArrayPlotData, Plot, Legend, cbrewer, palette11, palette14
from chaco.tools.api import DragZoom, PanTool, ZoomTool
from enable.api import ComponentEditor
from traits.api import *
from traitsui.api import *

def radar_range(p):
    return (p/5)**.25*180

def intercept_range(p):
    return (p/5)**.5*180

class SteamDuct(HasTraits):
    hd = Range(0, 80,40)
    plot = Instance(Plot)

    def update_plot(self, plot):
        h = np.linspace(0, 300., 301)
        plot.data.update_data(
            h=h,
            refractionrate=h,
        )

    def _hd_changed(self, event):
        self.update_plot(self.plot)

    def _plot_default(self):
        plot = Plot(ArrayPlotData())
        self.update_plot(plot)
        plot.x_axis.title = "refraction Rate [ Pa ]"
        plot.y_axis.title = "Height [ Pa ]"
        plot.orientation='v'
        plot.plot(("h", "refractionrate"), name='Refraction Rate', color='auto')
        legend = Legend()
        legend.plots = plot.plots
        plot.overlays.append(legend)
        return plot

    traits_view = View(
        Group(
            UCustom(
                "plot",
                editor=ComponentEditor(),
            ),
        ),
        Item('hd', label='波导高度', editor_args={'mode': 'slider'}),
    )


class Refraction_rate(HasTraits):
    temperatureRange = Range(-40, 60)
    plot = Instance(Plot)

    def update_plot(self, plot):
        p = np.arange(5000, 101325, 100)
        plot.data.update_data(
            p=p,
            N0=p,
        )

    def _temperatureRange_changed(self, event):
        self.update_plot(self.plot)

    def _plot_default(self):
        plot = Plot(ArrayPlotData())
        self.update_plot(plot)
        plot.title = 'refraction Rate'
        plot.x_axis.title = "Pressure [ Pa ]"
        plot.y_axis.title = "refraction Rate [ Pa ]"
        plot.plot(("p", "N0"), name='Humidity=0', color='auto')
        legend = Legend()
        legend.plots = plot.plots
        plot.overlays.append(legend)
        return plot

    traits_view = View(
        Group(
            UCustom(
                "plot",
                editor=ComponentEditor(),
            ),
        ),
        Item('temperatureRange', label='温度[℃]', editor_args={'mode': 'slider'}),
    )


class BasicRelation(HasTraits):
    plot = Instance(Plot)

    def _plot_default(self):
        plot = Plot(ArrayPlotData())
        # self.clear_plot(plot)
        return plot

    fun = Enum('截获因子', )
    def _fun_changed(self):
        plot = self.plot
        self.clear_plot(plot)
        match self.fun:
            case '截获因子':
                x = np.linspace(0, 50, 401)
                plot.data.update_data(
                    x=x,
                    y1=radar_range(x),
                    y2=intercept_range(x),
                )
                plot.x_axis.title = "Transmission Power (kW)"
                plot.y_axis.title = "Distance (km)"
                plot.plot(('x', 'y1'), name='Detection Distance', color='auto')
                plot.plot(('x', 'y2'), name='Intercept Distance', color='auto')
        legend = Legend()
        legend.plots = plot.plots
        plot.overlays.append(legend)
        plot.request_redraw()

    def traits_init(self):
        self._fun_changed()

    def clear_plot(self, plot):
        plots = [p for p in plot.plots.keys()]
        for p in plots:
            plot.delplot(p)
        for overlay in plot.overlays:
            plot.overlays.remove(overlay)
        plot.request_redraw()

    traits_view = View(
        Group(
            UCustom(
                "plot",
                editor=ComponentEditor(),
                resizable=True
            ),
            HGroup(
                UItem('fun'),
                spring,
            ),
        ),
        width=1000,
        height=800,
        resizable=True,
    )


class Form(HasTraits):
    basic_relation = Instance(BasicRelation, ())
    refraction_rate=Instance(Refraction_rate,())
    steamDuct=Instance(SteamDuct,())
    traits_view = View(
        Tabbed(
            UCustom('basic_relation', label='基本参数'),
            UCustom('refraction_rate', label='大气折射率',),
            UCustom('steamDuct', label='蒸发波导',),
        ),
        width=1000,
        height=800,
        resizable=True,
    )


if __name__ == '__main__':
    Form().configure_traits()
